// 只能处理1024内的排序
// #pragma once
#include <cstdio>
#include <iostream>
#include "myType.cuh"
#include "cuda_runtime_api.h"
#include <sys/time.h>
#include <typeinfo>
#include <algorithm>
#include <chrono>

#define GPU_CHECK(call)                               \
do                                                    \
{                                                     \
    const cudaError_t error_code = call;              \
    if (error_code != cudaSuccess)                    \
    {                                                 \
        printf("CUDA Error:\n");                      \
        printf("    File:       %s\n", __FILE__);     \
        printf("    Line:       %d\n", __LINE__);     \
        printf("    Error code: %d\n", error_code);   \
        printf("    Error text: %s\n",                \
            cudaGetErrorString(error_code));          \
        exit(1);                                      \
    }                                                 \
} while (0)

inline __device__ void gpu_swap(iAS* a, iAS* b) {
    iAS temp = *b;
    *b = *a;
    *a = temp;
}

void cpu_swap(iAS** a, iAS** b) {
    iAS* temp = *a;
    *a = *b;
    *b = temp;
}

// 大于比较符号
inline __device__  bool gpu_compare(const iAS& a, const iAS& b) {
    if(a.score != b.score) return a.score > b.score;
    return a.idx < b.idx;
}

bool cpu_compare(const iAS& a, const iAS& b) {
    if(a.score != b.score) return a.score > b.score;
    return a.idx < b.idx;
}

inline unsigned int cpu_dtNextPow2_p(unsigned int v) {
    unsigned cishu = 0, vv = 1;
    while(v > vv) {
        vv <<= 1;
        cishu++;
    }
	return cishu;
}

inline __device__ unsigned int gpu_dtNextPow2_p(unsigned int v) {
    unsigned cishu = 0, vv = 1;
    while(v > vv) {
        vv <<= 1;
        cishu++;
    }
	return cishu;
}

inline unsigned int cpu_dtNextPow2(unsigned int v)
{
	v--;
	v |= v >> 1;
	v |= v >> 2;
	v |= v >> 4;
	v |= v >> 8;
	v |= v >> 16;
	v++;
	return v;
}

inline __device__ unsigned int gpu_dtNextPow2(unsigned int v)
{
	v--;
	v |= v >> 1;
	v |= v >> 2;
	v |= v >> 4;
	v |= v >> 8;
	v |= v >> 16;
	v++;
	return v;
}

// 原版双调排序
__global__ void gBitonicSort(iAS* data, int n_p)
{
    unsigned int tid = threadIdx.x;

    register int stride_p, half_stride_p, s_p, hs_p, hs, i, j, k, hn;
    register bool orange;
    hn = 1 << (n_p - 1);
    half_stride_p = 0;
    
    for (stride_p = 1; stride_p <= n_p; stride_p++) {
        s_p = stride_p;
        while (s_p >= 1) {
            hs_p = s_p - 1;
            hs = 1 << hs_p;
            for (i = tid; i < hn; i += blockDim.x) {
                orange = (i >> half_stride_p) % 2 == 0;
                j = ((i >> hs_p) << s_p) + (i % hs);
                k = j + hs;
                if ((orange && !gpu_compare(data[j], data[k])) || (!orange && gpu_compare(data[j], data[k]))) {
                    gpu_swap(&data[k], &data[j]);
                }
            }
            __syncthreads();
            s_p = hs_p;
        }
        half_stride_p++;
    }
}

// 双调排序 n_p 可以等于0 但是那样的话不会排序，代表2^0=1
__device__ void gpu_topk_impl(iAS *data, const int n_p) {
    unsigned int tid = threadIdx.x;

    register int stride_p, half_stride_p, s_p, hs_p, hs, i, j, k, hn;
    register bool orange;
    hn = 1 << (n_p - 1);
    half_stride_p = 0;
    
    for (stride_p = 1; stride_p <= n_p; stride_p++) {
        s_p = stride_p;
        while (s_p >= 1) {
            hs_p = s_p - 1;
            hs = 1 << hs_p;
            for (i = tid; i < hn; i += blockDim.x) {
                orange = (i >> half_stride_p) % 2 == 0;
                j = ((i >> hs_p) << s_p) + (i % hs);
                k = j + hs;
                if ((orange && !gpu_compare(data[j], data[k])) || (!orange && gpu_compare(data[j], data[k]))) {
                    gpu_swap(&data[k], &data[j]);
                }
            }
            __syncthreads();
            s_p = hs_p;
        }
        half_stride_p++;
    }
}

__global__ void gpu_topk(iAS *data, iAS *d_result, const unsigned int length) {
    register auto tid = blockIdx.x * blockDim.x + threadIdx.x;
    
    // printf("adasd\n");

    if(tid >= length) {
        return;
    }

    __shared__ iAS temp_result[1024];

    // 赋值没有问题
    temp_result[threadIdx.x].score = data[tid].score; 
    temp_result[threadIdx.x].idx = data[tid].idx;
    
    if(threadIdx.x > 0) {
        return;
    }

    register unsigned int temp_result_length = 0;
    register unsigned int n_p = 0;
    register unsigned int n = 0;
    if(threadIdx.x == 0) {
        temp_result_length = (blockDim.x <= length - tid ? blockDim.x : length - tid);
        n_p = gpu_dtNextPow2_p(temp_result_length);
        n = gpu_dtNextPow2(temp_result_length);
        for(int i = temp_result_length; i < n; i++) temp_result[i].score = -1;
    }
    __syncthreads();

    gpu_topk_impl(temp_result, n_p);

    // printf("adasd\n");

    if(threadIdx.x > 0) {
        return;
    }

    printf("%d,%d,%d\n",temp_result_length, n_p, n);

    // 赋值到变量
    register int offset = blockIdx.x * TOPK;
    register int small_len = temp_result_length < TOPK ? temp_result_length : TOPK;
#pragma unroll
    for(register auto i = 0; i < small_len; i++) {
        printf("%d %f\n", temp_result[i].idx, temp_result[i].score);
        // printf("%d", i + blockIdx.x * TOPK);
        d_result[i + offset].score = temp_result[i].score;
        d_result[i + offset].idx = temp_result[i].idx;
        // printf("%d %f\n", d_result[i + blockIdx.x * TOPK].idx, d_result[i + blockIdx.x * TOPK].score);        
    }
}


// Initialize data on the host.
void initialize_data(iAS *dst, unsigned int nitems) {
    // Fixed seed for illustration
    srand(2047);

    // Fill dst with random values
    for (unsigned i = 0; i < nitems; i++) {
        dst[i].idx = i;
        dst[i].score = rand() % nitems;
    }
}

// Verify the results.
void print_results(iAS *results_d, int n) {
    iAS *results_h = new iAS[n];
    GPU_CHECK(cudaMemcpy(results_h, results_d, n * sizeof(iAS), cudaMemcpyDeviceToHost));
    std::cout << "Sort data : \n";
    for (int i = 0; i < n; ++i){
        std::cout << results_h[i].idx << " " << results_h[i].score << " " << "\n";
    }
    std::cout << std::endl;
    delete[] results_h;
}

void check_results(iAS *results_d, iAS *data_h, int n)
{
    iAS *results_h = new iAS[n];
    GPU_CHECK(cudaMemcpy(results_h, results_d, n*sizeof(iAS), cudaMemcpyDeviceToHost));

    for (int i = 0; i < n ; ++i) {
        if(results_h[i].idx != data_h[i].idx) {
            std::cout << "Invalid item[" << i << "]: " << results_h[i].idx << " greater than " << data_h[i].idx << std::endl;
            std::cout << "Invalid item[" << i << "]: " << results_h[i].score << " greater than " << data_h[i].score << std::endl;
            // exit(EXIT_FAILURE);
        }
    }
        // if (cpu_compare(results_h[i], results_h[i-1]))
        // {
        //     std::cout << "Invalid item[" << i-1 << "]: " << results_h[i-1].idx << " greater than " << results_h[i].idx << std::endl;
        //     std::cout << "Invalid item[" << i-1 << "]: " << results_h[i-1].score << " greater than " << results_h[i].score << std::endl;
        //     exit(EXIT_FAILURE);
        // }

    std::cout << "OK" << std::endl;
    delete[] results_h;
}

// Main entry point.
int main(int argc, char **argv) {
    // Find/set device and get device properties
    int device = 0;
    cudaDeviceProp deviceProp;
    GPU_CHECK(cudaGetDeviceProperties(&deviceProp, device));
    
    if (!(deviceProp.major > 3 ||
          (deviceProp.major == 3 && deviceProp.minor >= 5))) {
        printf("GPU %d - %s  does not support CUDA Dynamic Parallelism\n Exiting.",
            device, deviceProp.name);
        return 0;
    }

    int num_items = atoi(argv[1]); // 7853052
    unsigned int n = cpu_dtNextPow2(num_items);
    bool verbose = num_items <= 32;


    // Create input data
    iAS *h_data = 0;
    iAS *d_data = 0;
    iAS *d_result = 0;
    iAS *d_mid_result = 0;

    // Allocate CPU memory and initialize data.
    std::cout << "Initializing data:" << std::endl;
    h_data = (iAS *)malloc(n * sizeof(iAS));
    initialize_data(h_data, num_items);
    if (verbose) {
       std::cout << "Raw  data : \n";
        for (int i = 0; i < num_items; i++)
            std::cout << h_data[i].idx << " " << h_data[i].score << "\n";
    }

    for(int i = num_items; i < n; i++) {h_data[i].score = -1;}
    num_items = n;

    // Allocate GPU memory.
    GPU_CHECK(cudaMalloc((void **)&d_data, num_items * sizeof(iAS)));
    GPU_CHECK(cudaMalloc((void **)&d_result, num_items * sizeof(iAS)));
    GPU_CHECK(cudaMalloc((void **)&d_mid_result, num_items * sizeof(iAS)));
    
    // GPU_CHECK(cudaMemset(d_mid_result, -1, num_items * sizeof(iAS)));
    // GPU_CHECK(cudaMemset(d_result, -1, num_items * sizeof(iAS)));

    GPU_CHECK(cudaMemcpy(d_data, h_data, num_items * sizeof(iAS), cudaMemcpyHostToDevice));

    
    // Execute
    // printf("asdas");

    std::cout << "Running quicksort on " << num_items << " elements" << std::endl;
    
    // printf("asdas");

    std::chrono::high_resolution_clock::time_point t1 = std::chrono::high_resolution_clock::now();
    // printf("asdas");

    //grid要不大于10

    // 第一轮 grid = 7669
    int block = 512;
    int temp_grid = (num_items + block - 1) / block;
    unsigned int temp_length = num_items;
    // printf("%d %d", num_items, dtNextPow2_p(num_items));

    // gpu_topk<<<temp_grid, block>>>(d_data, d_mid_result, temp_length, false);
    // cudaDeviceSynchronize();
    // temp_length = (temp_grid - 1) * TOPK + (temp_length % block < TOPK ? temp_length % block : TOPK);
    // printf("a%d\n", temp_length);
    // print_results(d_mid_result, temp_length);


    // while(temp_length > 1024) {
    //     temp_grid = (temp_length + block - 1) / block;
    //     cpu_swap(&d_mid_result, &d_result);
    //     // print_results(d_result, temp_length);
    //     gpu_topk<<<temp_grid, block>>>(d_result, d_mid_result, temp_length, false);
    //     cudaDeviceSynchronize();

    //     temp_length = (temp_grid - 1) * TOPK + (temp_length % block < TOPK ? temp_length % block : TOPK);
    //     // printf("b%d\n", temp_length);
    //     // print_results(d_mid_result, temp_length);
    // }
    // printf("while finish\n");
    // register unsigned int temp_result_length = (blockDim.x <= length - tid ? blockDim.x : length - tid);
    // printf("asdas");
    unsigned int n_p = cpu_dtNextPow2_p(temp_length);
    
    // for(int i = temp_length; i < n; i++) d_data[i].score = -1;
    
    gBitonicSort<<<1, 1024>>>(d_data, n_p);
    // gpu_topk<<<1, temp_length>>>(d_data, d_result, temp_length);
    // cudaDeviceSynchronize();
    print_results(d_data, TOPK);

    
    std::chrono::high_resolution_clock::time_point t2 = std::chrono::high_resolution_clock::now();


    std::cout<<std::chrono::duration_cast<std::chrono::milliseconds>(t2-t1).count()<<std::endl;

    // print result
    // print_results(d_result, TOPK);
    // check result
    // std::cout << "Validating results: ";
    std::sort(h_data, h_data + num_items);
    check_results(d_data, h_data, TOPK);
    free(h_data);
    GPU_CHECK(cudaFree(d_data));

    return 0;
}

/* 
编译：nvcc -o quicksort_cuda --gpu-architecture=sm_70 -rdc=truequicksort_cuda.cu
*/